from __future__ import annotations
from typing import Type, Dict, Optional, Tuple, List, Union
import peewee
from peewee import CharField, IntegerField, DateTimeField, BooleanField, fn


from common.geometry.geom import Geom

from catalog.models.base_vector_item import BaseVectorItem, extract_base_info


class PlaceRecord(BaseVectorItem):

    collection_id = "PAII/PlaceV2Test"
    name_ori = CharField()
    name_en = CharField()
    # region_level: str[List[Integer]], assuming a place can have more than 1 region level.
    # "[1,2,3]" represent this place belongs to 3 region levels.
    region_level = CharField()
    # region_level_bin: Binary integer to represent the region level
    # From right to left, each binary digit represent if the place belong to this level.
    # 10: region level 1.
    # 1100: region level 2 and 3.
    region_level_bin = IntegerField()

    parent_name_ori = CharField()
    parent_name_en = CharField()
    gid = CharField()
    data_published_time = DateTimeField()
    data_source = CharField()
    year = IntegerField(index=True)
    latest_data = BooleanField(index=True)  # If the data is latest year
    comments = CharField()

    @classmethod
    def __query_with_name(
        cls: Type[PlaceRecord],
        query: peewee.Query,
        name_column: CharField,
        target_name: str,
        name_exact_mode: bool = False,
    ):
        name_column = fn.LOWER(name_column)
        target_name = target_name.lower()
        if name_exact_mode:
            if isinstance(target_name, list):
                query = query.where(name_column.in_(target_name))
            else:
                query = query.where(name_column == target_name)
        else:
            if isinstance(target_name, list):
                term = name_column.contains(target_name[0])
                if len(target_name) > 1:
                    for name_tmp in target_name[1:]:
                        term = term | name_column.contains(name_tmp)
                query = query.where(term)
            else:
                query = query.where(name_column.contains(target_name))

        return query

    @classmethod
    def __check_if_only_english(cls: Type[PlaceRecord], input_str: str):
        try:
            input_str.encode(encoding="utf-8").decode("ascii")
        except UnicodeDecodeError:
            return False
        else:
            return True

    @classmethod
    def create_one_item(cls: Type[PlaceRecord], info_item: Dict):
        base_info = extract_base_info(info_item)
        properties = info_item["properties"]

        return cls(
            **properties,
            **base_info,
        )

    @classmethod
    def region_level_list_to_binary(cls, region_level_list):
        max_region_level = 10
        assert max(region_level_list) < max_region_level
        region_flag_list = ["0" for _ in range(max_region_level)]
        for val in region_level_list:
            region_flag_list[val] = "1"

        region_bin_str = "".join(region_flag_list[::-1])
        return int(region_bin_str, 2)

    @classmethod
    def binary_to_region_level_list(cls, region_binary):
        binary_str = bin(region_binary)[2:][::-1]
        region_list = []
        for idx, val in binary_str:
            if val == "1":
                region_list.append(idx)

        return region_list

    @classmethod  # noqa C901
    def construct_query(
        cls: Type[PlaceRecord],
        spatial: Optional[Geom] = None,
        include_invalidated: bool = False,
        target_name: Optional[str] = None,
        target_name_exact_mode: bool = False,
        region_level: Optional[Union[int, List[int]]] = None,
        target_parent_name: Optional[str] = None,
        target_parent_name_exact_mode: bool = False,
        year: Optional[str] = "latest",
    ) -> Tuple[peewee.Query, int]:
        """
        Construct Query function. Do not use directly. Use query_many_items instead.
        """
        query, n_conditions = super(PlaceRecord, cls).construct_query(
            spatial=spatial, include_invalidated=include_invalidated
        )

        if year is not None:
            if year == "latest":
                query = query.where(cls.latest_data)
            else:
                query = query.where(cls.year == int(year))
            n_conditions += 1

        if target_name is not None:
            target_name = target_name.lower()
            all_english_flag = cls.__check_if_only_english(target_name)
            check_column = cls.name_en if all_english_flag else cls.name_ori
            query = cls.__query_with_name(
                query, check_column, target_name, target_name_exact_mode
            )
            n_conditions += 1

        if region_level is not None:
            if isinstance(region_level, int):
                region_level = [region_level]

            region_level_bin = cls.region_level_list_to_binary(region_level)
            query = query.where(cls.region_level_bin.bin_and(region_level_bin) > 0)
            n_conditions += 1

        if target_parent_name is not None:
            assert isinstance(target_parent_name, str)
            target_parent_name = target_parent_name.lower()
            all_english_flag = cls.__check_if_only_english(target_parent_name)
            check_column = (
                cls.parent_name_en if all_english_flag else cls.parent_name_ori
            )

            query = cls.__query_with_name(
                query, check_column, target_parent_name, target_parent_name_exact_mode
            )

            n_conditions += 1

        return query, n_conditions

    @classmethod
    @BaseVectorItem.auto_connection_value
    def query_many_items(
        cls: Type[PlaceRecord],
        spatial: Optional[Geom] = None,
        include_invalidated: bool = False,
        target_name: Optional[str] = None,
        target_name_exact_mode: bool = False,
        region_level: Optional[Union[int, List[int]]] = None,
        target_parent_name: Optional[str] = None,
        target_parent_name_exact_mode: bool = False,
        year: Union[int, str, None] = "latest",
    ) -> List[PlaceRecord]:
        """
        Query PlaceRecord results based on arguments.
        A list of Records will be returned

        Parameters
        ----------
        spatial : Optional[Geom], optional
            Find all PlaceRecords intersects with input geometry, by default None
        include_invalidated : bool, optional
            Whether to include invalidated records, by default False
        target_name : Optional[str], optional
            The name of AOI, can be both in Chinese or English,
            To place record in China, please input Chinese characters such as 北京市.
            For all other places in other countries, please use English.
            For example, Santa Clara, the name is case sensitive.
            By default None
        target_name_exact_mode : bool, optional
            Whether to use the exact mode in target location name.
            If True, name must be an exact match, otherwise,
            if the given name is within the field's value then it is a match.
        region_level : Optional[Union[int, List[int]]], optional
            The region level of the AOI:
                0: Country,
                1: Province / state,
                2: Prefecture (Only for areas in China)
                3: County
            by default None
        target_parent_name : Optional[str], optional
            Name of AOI parent, all areas' parent name contains target_parent_name
            will be returned unless target_parent_name_exact_mode is True
            For example, California is the parent of Santa Clara, United States
            is the parent of both California and Santa Clara. 辽宁省 is the parent of 沈阳市,
            中华人民共和国 is the parent of both 辽宁省 and 沈阳市
            By default None
        target_parent_name_exact_mode : bool, optional
            Whether to use the exact mode in target parent location name.
            If True, name must be an exact match, otherwise,
            if the given name is within the field's value then it is a match.
            By default False
        year : str, optional
            Which year of data should be selected. Currently, for China dataset, we have year from 2014 to 2022.
            Source from shengshixian.com
            For United States dataset, we have year from 2018 to 2021, source from census.gov.
            If set to int or a digit string, the data of selected year will be returned.
            If set to None, the dataset with all years will be returned.
            If set to "latest", only the data for latest year will be returned. for example,
            areas in China of year 2022 will be returned and areas in USA of year 2021 will be returned.

        Returns
        -------
        result_list : List[PlaceRecord]
            The PlaceRecord result list.
        """
        year = str(year).lower() if year is not None else None

        query, _ = cls.construct_query(
            spatial=spatial,
            include_invalidated=include_invalidated,
            target_name=target_name,
            target_name_exact_mode=target_name_exact_mode,
            region_level=region_level,
            target_parent_name=target_parent_name,
            target_parent_name_exact_mode=target_parent_name_exact_mode,
            year=year,
        )
        result_list: List[PlaceRecord] = [item for item in query]

        return result_list
